package com.cg.ai.guardrail;

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.guardrail.InputGuardrail;
import dev.langchain4j.guardrail.InputGuardrailResult;

import java.util.Set;

/**
 * 安全输入护轨
 * 用于检测和过滤用户输入中的敏感内容
 */
public class SafeInputGuardrail implements InputGuardrail {

    /**
     * 敏感词汇集合
     * 可根据实际需求扩展
     */
    private static final Set<String> SENSITIVE_WORDS = Set.of(
        "kill", "evil", "hack", "attack", "virus", "malware",
        "暴力", "攻击", "病毒", "恶意", "破坏", "非法"
    );

    /**
     * 验证用户输入是否安全
     * @param userMessage 用户消息
     * @return 验证结果
     */
    @Override
    public InputGuardrailResult validate(UserMessage userMessage) {
        try {
            // 获取用户输入并转换为小写以确保大小写不敏感
            String inputText = userMessage.singleText().toLowerCase();
            
            // 使用正则表达式分割输入文本为单词
            String[] words = inputText.split("\\W+");
            
            // 遍历所有单词，检查是否存在敏感词
            for (String word : words) {
                if (SENSITIVE_WORDS.contains(word)) {
                    return fatal("检测到敏感词汇: " + word + "，请使用文明用语进行咨询。");
                }
            }
            
            // 检查是否包含过长的输入（防止恶意输入）
            // 放宽限制，适配面试场景需要更长上下文
            if (inputText.length() > 10000) {
                return fatal("输入内容过长，请控制在10000字符以内。");
            }
            
            return success();
        } catch (Exception e) {
            // 如果验证过程出现异常，为了安全起见，拒绝输入
            return fatal("输入验证失败，请重新输入。");
        }
    }
}

